from hyperopt import fmin, tpe, hp, Trials, STATUS_OK
from hyperopt.pyll import scope    
from sklearn.ensemble import RandomForestClassifier    
from sklearn.model_selection import train_test_split, cross_validate
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
import numpy as np  
import pandas as pd  
import mlflow 
import json 
from mlflow.tracking import MlflowClient   
from xgboost import XGBClassifier
import lightgbm as lgb
from catboost import CatBoostClassifier

client = MlflowClient() 

def calculate_metrics_clf(y_true, y_pred, y_pred_proba,class_type):
    """计算并返回分类任务的评估指标"""
    accuracy = accuracy_score(y_true, y_pred)
    
    # 根据分类类型设置参数
    average = 'binary' if class_type == 'binary' else 'macro'
    
    precision = precision_score(y_true, y_pred, average=average) 
    recall = recall_score(y_true, y_pred, average=average)
    f1 = f1_score(y_true, y_pred, average=average)
    
    # 只有二分类计算AUC
    auc = None
    if class_type == 'binary' and y_pred_proba is not None:
        auc = roc_auc_score(y_true, y_pred_proba[:, 1])
        
    return accuracy, precision, recall, f1, auc

##================================ 随机森林超参数搜索 ============================ ##
def rf_clf_search_space(n_min, n_max, leaf_min, leaf_max,
                    d_min, d_max, criterion_list,
                    max_features_list):

    assert n_min <= n_max 
    assert leaf_min <= leaf_max 
    assert d_min <= d_max
    
    search_space = {   
        'n_estimators': scope.int(hp.quniform('n_estimators', n_min, n_max, 1)),
        'min_samples_leaf': scope.int(hp.quniform('min_samples_leaf', leaf_min, leaf_max, 1)), 
        'max_features': hp.choice('max_features', max_features_list + [None]),  
        'max_depth': scope.int(hp.quniform('max_depth', d_min, d_max, 1)),
        'criterion': hp.choice('criterion', criterion_list)
    } 

    return search_space


def cal_fmin_rf_clf(X_train, y_train, search_space, experiment_name, run_name, 
                    X_val=None, y_val=None, max_evals=20, cv=5, class_type='binary'):   
    """
    class_type: 'binary' 或 'multiclass'，指定分类类型 
    """
    mlflow.sklearn.autolog() 
    
    # 设置MLflow实验 
    try:
        mlflow.create_experiment(experiment_name)
    except mlflow.exceptions.MlflowException:
        pass
    mlflow.set_experiment(experiment_name)
    

    def objective(params: dict): 
        run_name_tmp = f"{run_name}_trial_{len(trials.trials)}"  
        with mlflow.start_run(nested=True,run_name=run_name_tmp):  
            model = RandomForestClassifier(** params, random_state=42, n_jobs=8)
            
            if cv > 1:
                # 交叉验证评估
                scoring = {
                    'accuracy': 'accuracy',
                    'precision': 'precision_macro' if class_type == 'multiclass' else 'precision',
                    'recall': 'recall_macro' if class_type == 'multiclass' else 'recall',
                    'f1': 'f1_macro' if class_type == 'multiclass' else 'f1'
                }
                
                # 二分类额外添加AUC 
                if class_type == 'binary':
                    scoring['auc'] = 'roc_auc'  
                    
                results = cross_validate(model, X_train, y_train, cv=cv, scoring=scoring) 
                
                accuracy = np.mean(results['test_accuracy'])
                precision = np.mean(results['test_precision'])
                recall = np.mean(results['test_recall'])
                f1 = np.mean(results['test_f1'])
                auc = np.mean(results['test_auc']) if class_type == 'binary' else None
                
                # 拟合完整模型
                model.fit(X_train, y_train)  
                
            else: 
                # 验证集评估
                model.fit(X_train, y_train)
                y_pred = model.predict(X_val)
                y_pred_proba = model.predict_proba(X_val)
                accuracy, precision, recall, f1, auc = calculate_metrics_clf(y_val, y_pred, y_pred_proba,class_type)
                
                # 评估之后，使用所有数据训练
                model.fit(pd.concat([X_train, X_val], axis=0), pd.concat([y_train, y_val], axis=0))  
            
            # 记录指标
            metrics = {
                "val_accuracy": accuracy,
                "val_precision": precision,
                "val_recall": recall,
                "val_f1": f1,
                "val_loss": 1-auc if  class_type == 'binary' else 1-f1  

            }
            
            if class_type == 'binary' and auc is not None:
                metrics["val_auc"] = auc
                
            mlflow.log_metrics(metrics)

            params['features'] = json.dumps(X_train.columns.tolist() )
            params['target'] = y_train.name
            params['task_type'] = class_type  
            mlflow.log_params(params) 
            
        # 以1 - accuracy作为损失函数（因为hyperopt是最小化损失）
        return {'loss': 1-auc if   class_type == 'binary' else 1-f1  , 'status': STATUS_OK}
    
    # 数据处理逻辑
    if cv > 1 and (X_val is not None and y_val is not None):
        X_train = pd.concat([X_train, X_val], axis=0)
        y_train = pd.concat([y_train, y_val], axis=0) 
    elif cv <= 1 and (X_val is None or y_val is None):
        X_train, X_val, y_train, y_val = train_test_split(
            X_train, y_train, test_size=0.2, random_state=42
        )

    with mlflow.start_run(run_name=run_name,nested=True) as run :
        trials = Trials()
        best_params = fmin(  
            fn=objective,
            space=search_space,
            algo=tpe.suggest,
            max_evals=max_evals,
            trials=trials
        ) 
    
        run_id = run.info.run_id 
        run = mlflow.get_run(run_id)
        experiment_id = run.info.experiment_id  

        best_run = client.search_runs(
                experiment_ids=experiment_id,
                filter_string=f"tags.mlflow.parentRunId = '{run_id}'",
                order_by=['metrics.val_f1','metrics.val_recall DESC', 'start_time DESC'],
                max_results=1,
                )[0]


        run_results ={
            'metrics':best_run.data.metrics ,
            'params' :  best_run.data.params  
        }
    
    return run_results  


##================================ XGBOOST超参数搜索 ============================ ## 

def xgb_clf_search_space(n_min, n_max, d_min, d_max,
                         lr_min, lr_max, gamma_min, gamma_max,
                         ss_min, ss_max, csbt_min, csbt_max, reg_lambda_min, reg_lambda_max,
                         class_type='binary'):

    assert n_min <= n_max
    assert d_min <= d_max
    assert lr_min <= lr_max
    assert gamma_min <= gamma_max
    assert ss_min <= ss_max
    assert csbt_min <= csbt_max
    assert reg_lambda_min <= reg_lambda_max

    # 根据分类类型设置目标函数
    objective = 'binary:logistic' if class_type == 'binary' else 'multi:softmax'
    
    search_space = {
        'n_estimators': scope.int(hp.quniform('n_estimators', n_min, n_max, 1)),
        'max_depth': scope.int(hp.quniform('max_depth', d_min, d_max, 1)),
        'learning_rate': hp.loguniform('learning_rate', np.log(lr_min), np.log(lr_max)),
        'gamma': hp.uniform('gamma', gamma_min, gamma_max),
        'subsample': hp.uniform('subsample', ss_min, ss_max),
        'colsample_bytree': hp.uniform('colsample_bytree', csbt_min, csbt_max),
        'reg_lambda': hp.uniform('reg_lambda', reg_lambda_min, reg_lambda_max),
        'objective': objective
    }
    
    # 多分类需要指定类别数量
    if class_type == 'multiclass':
        search_space['num_class'] = scope.int(hp.quniform('num_class', 3, 10, 1))  # 假设最多10个类别

    return search_space  


def cal_fmin_xgb_clf(X_train, y_train, search_space, experiment_name, run_name, 
                    X_val=None, y_val=None, max_evals=20, cv=5, class_type='binary'):
    
    
    mlflow.xgboost.autolog()
    
    try:
        mlflow.create_experiment(experiment_name)
    except mlflow.exceptions.MlflowException:
        pass
    mlflow.set_experiment(experiment_name)
    
    
    def objective(params: dict):
        run_name_tmp = f"{run_name}_trial_{len(trials.trials)}"  
        with mlflow.start_run(nested=True,run_name=run_name_tmp):  
            # 训练XGBoost分类模型
            model = XGBClassifier(**params, random_state=42, n_jobs=8)
            
            if cv > 1:
                # 交叉验证评估
                scoring = {
                    'accuracy': 'accuracy',
                    'precision': 'precision_macro' if class_type == 'multiclass' else 'precision',
                    'recall': 'recall_macro' if class_type == 'multiclass' else 'recall',
                    'f1': 'f1_macro' if class_type == 'multiclass' else 'f1'
                }
                
                if class_type == 'binary':
                    scoring['auc'] = 'roc_auc'
                    
                results = cross_validate(model, X_train, y_train, cv=cv, scoring=scoring)
                
                accuracy = np.mean(results['test_accuracy'])
                precision = np.mean(results['test_precision'])
                recall = np.mean(results['test_recall'])
                f1 = np.mean(results['test_f1'])
                auc = np.mean(results['test_auc']) if class_type == 'binary' else None
                
                # 交叉验证后拟合完整模型
                model.fit(X_train, y_train) 

            else:
                model.fit(X_train, y_train)
                y_pred = model.predict(X_val)
                y_pred_proba = model.predict_proba(X_val)
                accuracy, precision, recall, f1, auc = calculate_metrics_clf(y_val, y_pred, y_pred_proba,class_type)
                
                # 评估之后，使用所有数据训练
                model.fit(pd.concat([X_train, X_val], axis=0), pd.concat([y_train, y_val], axis=0)) 
            
            
            # 记录评估指标
            metrics = {
                "val_accuracy": accuracy,
                "val_precision": precision,
                "val_recall": recall,
                "val_f1": f1,
                "val_loss": 1-auc if   class_type == 'binary' else 1-f1  

            }
            
            if class_type == 'binary' and auc is not None:
                metrics["val_auc"] = auc 
                
            mlflow.log_metrics(metrics)
            
            # 记录XGBoost模型
            mlflow.xgboost.log_model(
                xgb_model=model,
                artifact_path="model",
                signature=mlflow.models.infer_signature(X_train, model.predict(X_train))
            )
            
            params['features'] = json.dumps(X_train.columns.tolist() )
            params['target'] = y_train.name
            params['task_type'] = class_type 
            mlflow.log_params(params) 
        # 以1 - accuracy作为损失函数
        return {'loss': 1-auc if   class_type == 'binary' else 1-f1 , 'status': STATUS_OK}
    
    
    if cv > 1 and (X_val is not None and y_val is not None):
        X_train = pd.concat([X_train, X_val], axis=0)
        y_train = pd.concat([y_train, y_val], axis=0)
    elif cv <= 1 and (X_val is None or y_val is None):
        X_train, X_val, y_train, y_val = train_test_split(
            X_train, y_train, test_size=0.2, random_state=42
        )
    
    with mlflow.start_run(run_name=run_name,nested=True) as run : 
        trials = Trials()
        best_params = fmin(
            fn=objective,
            space=search_space,
            algo=tpe.suggest,
            max_evals=max_evals,
            trials=trials
        ) 
    
        run_id = run.info.run_id 
        run = mlflow.get_run(run_id)
        experiment_id = run.info.experiment_id  

        best_run = client.search_runs(
                experiment_ids=experiment_id,
                filter_string=f"tags.mlflow.parentRunId = '{run_id}'",
                order_by=['metrics.val_loss','metrics.val_recall DESC', 'start_time DESC'],
                max_results=1,
                )[0]


        run_results ={
            'metrics':best_run.data.metrics ,
            'params' :  best_run.data.params  
        }
    
    return run_results  


##================================ LightGBM 超参数搜索 ============================ ##  
def lgb_clf_search_space(n_min, n_max, d_min, d_max,
                         lr_min, lr_max, leaves_min, leaves_max,
                         min_child_samples_min, min_child_samples_max,
                         ss_min, ss_max, csbt_min, csbt_max, reg_lambda_min, reg_lambda_max,
                         class_type='binary'):

    assert n_min <= n_max
    assert d_min <= d_max
    assert lr_min <= lr_max
    assert leaves_min <= leaves_max
    assert min_child_samples_min <= min_child_samples_max
    assert ss_min <= ss_max
    assert csbt_min <= csbt_max
    assert reg_lambda_min <= reg_lambda_max

    # 根据分类类型设置目标函数
    objective = 'binary' if class_type == 'binary' else 'multiclass'
    
    search_space = {
        'n_estimators': scope.int(hp.quniform('n_estimators', n_min, n_max, 1)),
        'max_depth': scope.int(hp.quniform('max_depth', d_min, d_max, 1)),
        'learning_rate': hp.loguniform('learning_rate', np.log(lr_min), np.log(lr_max)),
        'num_leaves': scope.int(hp.quniform('num_leaves', leaves_min, leaves_max, 1)),
        'min_child_samples': scope.int(hp.quniform('min_child_samples', min_child_samples_min, min_child_samples_max, 1)),
        'subsample': hp.uniform('subsample', ss_min, ss_max),
        'colsample_bytree': hp.uniform('colsample_bytree', csbt_min, csbt_max),
        'reg_lambda': hp.uniform('reg_lambda', reg_lambda_min, reg_lambda_max),
        'objective': objective
    }
    
    # 多分类需要指定类别数量
    if class_type == 'multiclass':
        search_space['num_class'] = scope.int(hp.quniform('num_class', 3, 10, 1))  # 假设最多10个类别

    return search_space 


def cal_fmin_lgb_clf(X_train, y_train, search_space, experiment_name, run_name, 
                    X_val=None, y_val=None, max_evals=20, cv=5, class_type='binary'):
    """
    使用hyperopt为LightGBM分类模型进行超参数优化
    
    参数:
        X_train: 训练特征
        y_train: 训练目标
        search_space: hyperopt格式的搜索空间
        experiment_name: MLflow实验名称
        run_name: MLflow运行名称
        X_val: 验证特征（可选）
        y_val: 验证目标（可选）
        max_evals: 最大搜索次数
        cv: 交叉验证折数
        class_type: 'binary' 或 'multiclass'，指定分类类型
    """
    # 初始化MLflow
    mlflow.lightgbm.autolog()
    try:
        mlflow.create_experiment(experiment_name)
    except mlflow.exceptions.MlflowException:
        pass
    mlflow.set_experiment(experiment_name)
    
    # hyperopt目标函数
    def objective(params):
        # 转换整数参数
        int_params = ['num_leaves', 'max_depth', 'n_estimators', 'min_child_samples', 'num_class']
        for param in int_params:
            if param in params:
                params[param] = int(params[param])
        
        run_name_tmp = f"{run_name}_trial_{len(trials.trials)}"  
        with mlflow.start_run(nested=True,run_name=run_name_tmp):  
            # 初始化LightGBM模型
            model = lgb.LGBMClassifier(
                **params,
                random_state=42,
                n_jobs=8,
                verbose=-1
            )
            
            if cv > 1:
                # 交叉验证评估
                scoring = {
                    'accuracy': 'accuracy',
                    'precision': 'precision_macro' if class_type == 'multiclass' else 'precision',
                    'recall': 'recall_macro' if class_type == 'multiclass' else 'recall',
                    'f1': 'f1_macro' if class_type == 'multiclass' else 'f1'
                }
                
                if class_type == 'binary':
                    scoring['auc'] = 'roc_auc'
                    
                results = cross_validate(model, X_train, y_train, cv=cv, scoring=scoring)
                
                accuracy = np.mean(results['test_accuracy'])
                precision = np.mean(results['test_precision'])
                recall = np.mean(results['test_recall'])
                f1 = np.mean(results['test_f1'])
                auc = np.mean(results['test_auc']) if class_type == 'binary' else None
                
                # 全量数据拟合模型
                model.fit(X_train, y_train)
            else:
                # 验证集评估
                model.fit(X_train, y_train)
                y_pred = model.predict(X_val)
                y_pred_proba = model.predict_proba(X_val)
                accuracy, precision, recall, f1, auc = calculate_metrics_clf(y_val, y_pred, y_pred_proba,class_type) 
                
                # 评估之后，使用所有数据训练
                model.fit(pd.concat([X_train, X_val], axis=0), pd.concat([y_train, y_val], axis=0)) 

            
            # 记录参数和指标
            metrics = {
                "val_accuracy": accuracy,
                "val_precision": precision,
                "val_recall": recall,
                "val_f1": f1,
                "val_loss": 1-auc if   class_type == 'binary' else 1-f1  

            }
            
            if class_type == 'binary' and auc is not None:
                metrics["val_auc"] = auc 
                
            mlflow.log_metrics(metrics)
            
            # 记录模型
            mlflow.lightgbm.log_model(
                lgb_model=model,
                artifact_path="model",
                signature=mlflow.models.infer_signature(X_train, model.predict(X_train))
            )

            params['features'] = json.dumps(X_train.columns.tolist())
            params['target'] = y_train.name
            params['task_type'] = class_type
            mlflow.log_params(params)
        
        # 以1 - accuracy作为损失函数
        return {'loss': 1-auc if   class_type == 'binary' else 1-f1  , 'status': STATUS_OK}
    
    # 数据处理逻辑
    if cv > 1 and (X_val is not None and y_val is not None):
        X_train = pd.concat([X_train, X_val], axis=0)
        y_train = pd.concat([y_train, y_val], axis=0)
    elif cv <= 1 and (X_val is None or y_val is None):
        X_train, X_val, y_train, y_val = train_test_split(
            X_train, y_train, test_size=0.2, random_state=42
        )
    
    # 执行hyperopt搜索
    with mlflow.start_run(run_name=run_name,nested=True) as run : 
        trials = Trials()
        best_params = fmin(
            fn=objective,
            space=search_space,
            algo=tpe.suggest,
            max_evals=max_evals,
            trials=trials,
            rstate=np.random.default_rng(42)  # 固定随机种子
        )
        
        run_id = run.info.run_id 
        run = mlflow.get_run(run_id)
        experiment_id = run.info.experiment_id  

        best_run = client.search_runs(
                experiment_ids=experiment_id,
                filter_string=f"tags.mlflow.parentRunId = '{run_id}'",
                order_by=['metrics.val_loss','metrics.val_recall DESC', 'start_time DESC'],
                max_results=1,
                )[0]


        run_results ={
            'metrics':best_run.data.metrics ,
            'params' :  best_run.data.params  
        }
    
    return run_results  


##================================ CatBoost 超参数搜索 ============================ ##  
def cat_clf_search_space(n_min, n_max, d_min, d_max,
                         lr_min, lr_max, l2_reg_min, l2_reg_max,
                         class_type='binary'):
    
    assert n_min <= n_max
    assert d_min <= d_max
    assert lr_min <= lr_max
    assert l2_reg_min <= l2_reg_max

    # 根据分类类型设置评估指标
    eval_metric = 'AUC' if class_type == 'binary' else 'MultiClass'
    
    search_space = {
        'n_estimators': scope.int(hp.quniform('n_estimators', n_min, n_max, 1)),
        'max_depth': scope.int(hp.quniform('max_depth', d_min, d_max, 1)),
        'learning_rate': hp.loguniform('learning_rate', np.log(lr_min), np.log(lr_max)),
        'l2_leaf_reg': hp.uniform('l2_leaf_reg', l2_reg_min, l2_reg_max),
        'eval_metric': eval_metric,
        'random_strength': hp.uniform('random_strength', 0.1, 10),
        'loss_function': 'Logloss' if class_type == 'binary' else 'MultiClass'
    }
    
    # 多分类需要指定类别数量
    if class_type == 'multiclass':
        search_space['classes_count'] = scope.int(hp.quniform('classes_count', 3, 10, 1))  # 假设最多10个类别

    return search_space


def cal_fmin_catboost_clf(X_train, y_train, search_space, experiment_name, run_name, 
                         X_val=None, y_val=None, max_evals=20, cv=5, class_type='binary'):
    """
    使用hyperopt为CatBoost分类模型进行超参数优化
    
    参数:
        X_train: 训练特征
        y_train: 训练目标
        search_space: hyperopt格式的搜索空间
        experiment_name: MLflow实验名称
        run_name: MLflow运行名称
        X_val: 验证特征（可选）
        y_val: 验证目标（可选）
        max_evals: 最大搜索次数
        cv: 交叉验证折数
        class_type: 'binary' 或 'multiclass'，指定分类类型
    """
    # 初始化MLflow
    mlflow.sklearn.autolog()
     
    try:
        mlflow.create_experiment(experiment_name)
    except mlflow.exceptions.MlflowException:
        pass
    mlflow.set_experiment(experiment_name)
    
    # hyperopt目标函数
    def objective(params):
        # 转换整数参数
        int_params = ['max_depth', 'n_estimators', 'classes_count']
        for param in int_params:
            if param in params:
                params[param] = int(params[param])
        
        run_name_tmp = f"{run_name}_trial_{len(trials.trials)}"  
        with mlflow.start_run(nested=True,run_name=run_name_tmp):  
            # 初始化CatBoost模型
            model = CatBoostClassifier(
                **params,
                random_state=42,
                thread_count=8,
                verbose=0
            )
            
            if cv > 1:
                # 交叉验证评估
                scoring = {
                    'accuracy': 'accuracy',
                    'precision': 'precision_macro' if class_type == 'multiclass' else 'precision',
                    'recall': 'recall_macro' if class_type == 'multiclass' else 'recall',
                    'f1': 'f1_macro' if class_type == 'multiclass' else 'f1'
                }
                
                if class_type == 'binary':
                    scoring['auc'] = 'roc_auc'
                    
                results = cross_validate(model, X_train, y_train, cv=cv, scoring=scoring)
                
                accuracy = np.mean(results['test_accuracy'])
                precision = np.mean(results['test_precision'])
                recall = np.mean(results['test_recall'])
                f1 = np.mean(results['test_f1'])
                auc = np.mean(results['test_auc']) if class_type == 'binary' else None
                
                # 全量数据拟合模型
                model.fit(X_train, y_train)
            else:
                # 验证集评估
                model.fit(X_train, y_train)
                y_pred = model.predict(X_val)
                y_pred_proba = model.predict_proba(X_val)
                accuracy, precision, recall, f1, auc = calculate_metrics_clf(y_val, y_pred, y_pred_proba,class_type) 
                
                # 评估之后，使用所有数据训练
                model.fit(pd.concat([X_train, X_val], axis=0), pd.concat([y_train, y_val], axis=0))  
            
            # 记录参数和指标
            metrics = {
                "val_accuracy": accuracy,
                "val_precision": precision,
                "val_recall": recall,
                "val_f1": f1,
                "val_loss": 1-auc if   class_type == 'binary' else 1-f1  

            }
            
            if class_type == 'binary' and auc is not None:
                metrics["val_auc"] = auc  
                
            mlflow.log_metrics(metrics)
            
            # 记录模型
            mlflow.catboost.log_model( 
                cb_model=model,
                artifact_path="model",
                signature=mlflow.models.infer_signature(X_train, model.predict(X_train))
            )

            params['features'] = json.dumps(X_train.columns.tolist() )
            params['target'] = y_train.name
            params['task_type'] = class_type  
            mlflow.log_params(params) 
        
        # 以1 - accuracy作为损失函数
        return {'loss': 1-auc if   class_type == 'binary' else 1-f1  , 'status': STATUS_OK}
    
    # 数据处理逻辑
    if cv > 1 and (X_val is not None and y_val is not None):
        X_train = pd.concat([X_train, X_val], axis=0)
        y_train = pd.concat([y_train, y_val], axis=0)
    elif cv <= 1 and (X_val is None or y_val is None):
        X_train, X_val, y_train, y_val = train_test_split(
            X_train, y_train, test_size=0.2, random_state=42
        )
    
    # 执行hyperopt搜索
    with mlflow.start_run(run_name=run_name,nested=True) as run : 
        trials = Trials()
        best_params = fmin(
            fn=objective,
            space=search_space,
            algo=tpe.suggest,
            max_evals=max_evals,
            trials=trials,
            rstate=np.random.default_rng(42)  # 固定随机种子
        )
        
        run_id = run.info.run_id 
        run = mlflow.get_run(run_id)
        experiment_id = run.info.experiment_id  

        best_run = client.search_runs(
                experiment_ids=experiment_id,
                filter_string=f"tags.mlflow.parentRunId = '{run_id}'",
                order_by=['metrics.val_loss','metrics.val_recall DESC', 'start_time DESC'],
                max_results=1,
                )[0]


        run_results ={
            'metrics':best_run.data.metrics ,
            'params' :  best_run.data.params  
        }
    
    return run_results  